import tensorflow as tf
import collections
import os

# return list of words from an article
def _read_words(filename):
    with tf.gfile.GFile(filename, "r") as f:
        return f.read().replace("\n", "<eos>").split()

# return dict of {word: id} for words in an article
def _build_vocab(filename):
    # read word as list
    data = _read_words(filename)
    len_all = len(data)
    data = data[: len_all // 10]

    # count word frequency
    counter = collections.Counter(data)
    # return a new list
    # with decreasing frequency, lexicographically ascending order
    # as form of (word, frequency)
    count_pairs = sorted(counter.items(), key=lambda x: (-x[1], x[0]))

    # separate word from its frequency
    words, _ = list(zip(*count_pairs))
    # return a dict of {word: id}
    # where smaller id means high frequency
    word_to_id = dict(zip(words, range(len(words))))

    return word_to_id

# return list of words' ids as order in article if exist in dict
def _file_to_word_ids(filename, word_to_id):
    data = _read_words(filename)
    len_all = len(data)
    data = data[: len_all // 10]
    len_all = len(data)
    len_train = len_all // 10 * 9
    len_valid = len_all // 100 * 5
    train_data = data[: len_train]
    valid_data = data[len_train : len_train + len_valid]
    test_data  = data[len_train + len_valid :]
    return ([word_to_id[word] for word in train_data if word in word_to_id],
        [word_to_id[word] for word in valid_data if word in word_to_id],
        [word_to_id[word] for word in test_data if word in word_to_id])


def fetch_data(data_path=None):
    """The Text8 dataset comes from
    http://mattmahoney.net/dc/text8.zip
    Args:
        data_path: string path to the directory where text8.zip has
            been extracted.
    Returns:
        tuple (train_data, valid_data, test_data, vocabulary)
        where each of the data objects can be passed to PTBIterator.
    """

    data_path = os.path.join(data_path, "text8")

    # use train date to build vocab dict
    word_to_id = _build_vocab(data_path)
    # use vocab dict to id words in file
    train_data, valid_data, test_data = _file_to_word_ids(data_path, word_to_id)
    vocabulary = len(word_to_id)
    return train_data, valid_data, test_data, vocabulary
